"""
YOLO11 API服务
使用 Ultralytics YOLO 官方接口（简化版）
参考: https://docs.ultralytics.com/zh/usage/python/

示例:
    from ultralytics import YOLO
    model = YOLO("/Users/yanlp/models/yolo11n.pt")
    results = model("image.jpg")
"""
import os
import logging
import yaml
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
import uvicorn
from prometheus_client import Counter, Histogram, generate_latest, CONTENT_TYPE_LATEST
from starlette.responses import Response
from ultralytics import YOLO

from .inference import predict_from_base64

# 配置日志
logging.basicConfig(
    level=os.getenv("LOG_LEVEL", "INFO"),
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# Prometheus指标
request_counter = Counter(
    "yolo_requests_total",
    "Total number of requests",
    ["endpoint", "status"]
)

request_duration = Histogram(
    "yolo_request_duration_seconds",
    "Request duration in seconds",
    ["endpoint"]
)

# 全局变量
model = None
config = None


def load_config():
    """加载配置文件"""
    config_path = os.path.join(os.path.dirname(__file__), "../config/config.yaml")
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    global model, config
    
    # 启动时只验证模型目录
    logger.info("正在初始化YOLO11服务...")
    try:
        config = load_config()
        model_config = config.get('model', {})
        
        # 获取模型目录路径
        local_model_path = os.getenv("LOCAL_MODEL_PATH") or model_config.get('local_model_path')
        
        if not local_model_path:
            raise ValueError("必须配置 LOCAL_MODEL_PATH 环境变量或在 config.yaml 中设置 local_model_path")
        
        # 提取模型目录（/app/models）
        model_dir = os.path.dirname(local_model_path)
        logger.info(f"模型目录: {model_dir}")
        
        # 检查模型目录是否存在
        if not os.path.exists(model_dir):
            raise FileNotFoundError(f"模型目录不存在: {model_dir}")
        
        logger.info("✓ 模型目录验证通过")
        logger.info("YOLO11服务初始化完成（模型将在推理时动态加载）")
        
    except Exception as e:
        logger.error(f"服务初始化失败: {str(e)}")
        raise
    
    yield
    
    # 关闭时清理资源
    logger.info("正在清理资源...")


app = FastAPI(
    title="YOLO11 API",
    description="基于YOLO11的目标检测服务（使用官方Python接口）",
    version="1.0.0",
    lifespan=lifespan
)

# CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class HealthResponse(BaseModel):
    """健康检查响应"""
    status: str
    model_loaded: bool
    version: str
    model_path: str
    task: str


class ChatCompletionRequest(BaseModel):
    """YOLO推理请求（参考Ultralytics API格式）"""
    source: str  # base64编码的图像数据
    model: Optional[str] = "yolo11n.pt"  # 模型文件名
    imgsz: Optional[int] = 640
    conf: Optional[float] = 0.25
    iou: Optional[float] = 0.45


@app.get("/health", response_model=HealthResponse)
async def health_check():
    """健康检查"""
    model_config = config.get('model', {})
    local_model_path = os.getenv("LOCAL_MODEL_PATH") or model_config.get('local_model_path')
    model_dir = os.path.dirname(local_model_path) if local_model_path else "unknown"
    
    return {
        "status": "healthy",
        "model_loaded": True,  # 动态加载模式
        "version": "1.0.0",
        "model_path": model_dir,
        "task": "dynamic"  # 根据请求动态加载
    }


@app.get("/")
async def root():
    """根路径"""
    return {
        "service": "YOLO11 API",
        "model": "yolo11",
        "version": "1.0.0",
        "description": "目标检测、分割、分类、姿态估计等多任务支持（base64图像推理，NPU加速）",
        "endpoints": {
            "inference": "/yolo11",
            "health": "/health",
            "metrics": "/metrics"
        },
        "documentation": "https://docs.ultralytics.com/zh/usage/python/",
        "api_reference": "https://docs.ultralytics.com/zh/hub/inference-api/#curl"
    }


@app.get("/metrics")
async def metrics():
    """Prometheus指标"""
    return Response(content=generate_latest(), media_type=CONTENT_TYPE_LATEST)


@app.post("/yolo11")
async def yolo11_inference(
    request: ChatCompletionRequest
):
    """
    YOLO11推理接口（兼容Ultralytics API格式）
    
    使用官方YOLO接口（简化版）:
        from ultralytics import YOLO
        model = YOLO("yolo11n.pt")
        results = model("image.jpg")
    
    参考: https://docs.ultralytics.com/zh/hub/inference-api/#curl
    
    请求体:
        {
            "source": "/9j/4AAQSkZJRg...",  // base64编码的图像数据
            "model": "yolo11s-seg.pt",
            "imgsz": 640,
            "conf": 0.25,
            "iou": 0.45
        }
    
    Returns:
        Ultralytics格式的响应（与官方API兼容）
    """
    endpoint = "/yolo11"
    request_counter.labels(endpoint=endpoint, status="received").inc()
    
    try:
        with request_duration.labels(endpoint=endpoint).time():
            # 检查是否提供了图像
            if not request.source:
                raise HTTPException(
                    status_code=400, 
                    detail="必须提供 source (base64编码的图像数据)"
                )
            
            logger.info(f"接收推理请求: model={request.model}, imgsz={request.imgsz}, conf={request.conf}, iou={request.iou}")
            
            # 获取模型目录
            model_config = config.get('model', {})
            local_model_path = os.getenv("LOCAL_MODEL_PATH") or model_config.get('local_model_path')
            model_dir = os.path.dirname(local_model_path)
            
            # 构造完整的模型路径
            model_path = os.path.join(model_dir, request.model)
            
            # 检查模型文件是否存在
            if not os.path.exists(model_path):
                raise HTTPException(
                    status_code=404,
                    detail=f"模型不存在: {request.model}"
                )
            
            logger.info(f"使用模型: {model_path}")
            
            # 动态加载并推理（官方简化接口）
            result = predict_from_base64(
                model_path,
                request.source,
                imgsz=request.imgsz,
                conf=request.conf,
                iou=request.iou
            )
            
            request_counter.labels(endpoint=endpoint, status="success").inc()
            
            # 返回Ultralytics格式的响应
            return result
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"推理失败: {str(e)}", exc_info=True)
        request_counter.labels(endpoint=endpoint, status="error").inc()
        raise HTTPException(status_code=500, detail=str(e))




if __name__ == "__main__":
    port = int(os.getenv("PORT", 8000))
    
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=port,
        workers=1,
        log_level=os.getenv("LOG_LEVEL", "info").lower()
    )
